#include "TCPSocket.h"
#include "NetworkProtocol.h"
#include "Encrypt.h"
#include "EncryptTable.h"

TCPSocket::TCPSocket()
{ 
	memset(m_bufOutput, 0, sizeof(m_bufOutput));
	memset(m_bufInput, 0, sizeof(m_bufInput));
	memset(m_bufMakePacket, 0, sizeof(m_bufMakePacket));

	m_scktState = Socket_Disconnected;
	m_scktClient = INVALID_SOCKET;
	m_nBlockSeconds = BLOCKSECONDS;

	m_keyDecrypt = 0;
	m_bFirstDecode = true;
	m_keyEncrypt = 0;
	m_bFirstEncode = true;
	m_pProtocol = 0;
}

void TCPSocket::closeSocket()
{
	if (m_pProtocol)
	{
		delete m_pProtocol;
		m_pProtocol = 0;
	}
#ifdef WIN32
	closesocket(m_scktClient);
#else
	close(m_scktClient);
#endif
}

bool TCPSocket::create(const char* pszServerIP, int nServerPort, int tagid, NetworkProtocol* protocol,
											 int nBlockSec, bool bKeepAlive /*= FALSE*/)
{
	m_pProtocol = protocol;
	m_nBlockSeconds = nBlockSec;
	if (m_pProtocol)
	{
		m_pProtocol->setTCPSocket(this);
	}
	if(pszServerIP == 0) {
		return false;
	}

	m_scktClient = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if(m_scktClient == INVALID_SOCKET) {
		closeSocket();
		return false;
	}

	if(bKeepAlive)
	{
		int	optval=1;
		if(setsockopt(m_scktClient, SOL_SOCKET, SO_KEEPALIVE, (char *) &optval, sizeof(optval)))
		{
			closeSocket();
			return false;
		}
	}

#ifdef WIN32
	DWORD nMode = 1;
	int nRes = ioctlsocket(m_scktClient, FIONBIO, &nMode);
	if (nRes == SOCKET_ERROR) {
		closeSocket();
		return false;
	}
#else
	fcntl(m_scktClient, F_SETFL, O_NONBLOCK);
#endif

	sockaddr_in	addr_in;
	memset((void *)&addr_in, 0, sizeof(addr_in));
	addr_in.sin_family = AF_INET;
	addr_in.sin_port = htons(nServerPort);
	unsigned long serveraddr = inet_addr(pszServerIP);
	if(serveraddr == INADDR_NONE)    
	{
		struct hostent *host;
		host = gethostbyname(pszServerIP);
		if (host == 0)
		{
			closeSocket();
			return false;
		}
		addr_in.sin_addr = *((struct in_addr *)host->h_addr_list[0]);
	}
	else
	{
		addr_in.sin_addr.s_addr = serveraddr;
	}

	if(connect(m_scktClient, (sockaddr *)&addr_in, sizeof(addr_in)) == SOCKET_ERROR) 
	{
		if (hasError()) {
			closeSocket();
			return false;
		}
	}

	m_nInbufLen        = 0;
	m_nOutbufLen    = 0;
	m_nMakePacketLen = 0;

	struct linger so_linger;
	so_linger.l_onoff = 1;
	so_linger.l_linger = 500;
	setsockopt(m_scktClient, SOL_SOCKET, SO_LINGER, (const char*)&so_linger, sizeof(so_linger));
	m_tag = tagid;
	m_scktState = Socket_Connecting;

	return true;
}
bool TCPSocket::polling()
{
	if (m_scktState == Socket_Disconnected) return false;
	if (m_scktClient == INVALID_SOCKET) return false;

	fd_set readset, writeset, exceptset;
	struct timeval timeout_tv;
	int result;

	FD_ZERO(&readset);
	FD_ZERO(&writeset);
	FD_ZERO(&exceptset);

	FD_SET(m_scktClient, &readset);
	FD_SET(m_scktClient, &writeset);
	FD_SET(m_scktClient, &exceptset);

	timeout_tv.tv_sec = m_nBlockSeconds;
	timeout_tv.tv_usec = 0;


	result = select(m_scktClient + 1, &readset, &writeset, &exceptset, &timeout_tv);
	if (result==SOCKET_ERROR) return false;
	if (result == 0) return true;

	if (FD_ISSET(m_scktClient, &exceptset)) 
	{
		destroy();
	}

	if (FD_ISSET(m_scktClient, &readset))
	{
		doRead();
	}

	if (FD_ISSET(m_scktClient, &writeset)) 
	{
		if (m_scktState == Socket_Connecting)
		{
			m_scktState = Socket_Connected;
		}

		doWrite();
	}
	return true;
}

bool TCPSocket::hasError()
{
#ifdef WIN32
	int err = WSAGetLastError();
	if(err != WSAEWOULDBLOCK) {
#else
	int err = errno;
	if(err != EINPROGRESS && err != EAGAIN) {
#endif
		return true;
	}

	return false;
}

bool TCPSocket::doRead()
{
	if (m_scktClient == INVALID_SOCKET) return false;
	if (m_scktState != Socket_Connected) return false;
	int recv_res = recv(m_scktClient, &m_bufInput[m_nInbufLen], INBUFSIZE - m_nInbufLen, 0);
	if (recv_res == SOCKET_ERROR) return false;
	if (recv_res == 0) return true;
	if (m_bFirstDecode)
	{
		m_idxDecodeTable = (unsigned int)m_bufInput[m_nInbufLen] % 10;
		m_keyDecrypt = m_bufInput[m_nInbufLen];
		m_keyDecrypt = Decrypt(m_idxDecodeTable, m_keyDecrypt, (unsigned char*)&m_bufInput[m_nInbufLen + 1],
			(unsigned char*)&m_bufInput[m_nInbufLen + 1], recv_res-1);
		memmove(&m_bufInput[m_nInbufLen], &m_bufInput[m_nInbufLen + 2], recv_res -2);
		recv_res -= 2;
		m_bFirstDecode = false;
	}
	else
	{
		m_keyDecrypt = Decrypt(m_idxDecodeTable, m_keyDecrypt, (unsigned char*)&m_bufInput[m_nInbufLen],
			(unsigned char*)&m_bufInput[m_nInbufLen], recv_res);
	}
	m_nInbufLen += recv_res;
	if (m_pProtocol)
	{
		int consumed = m_pProtocol->consumePacket(m_bufInput, m_nInbufLen);
		while ((consumed > 0)&&(m_nInbufLen > 0))
		{
			memmove(m_bufInput, &m_bufInput[consumed], m_nInbufLen - consumed);
			m_nInbufLen -= consumed;
			consumed = m_pProtocol->consumePacket(m_bufInput, m_nInbufLen);
		}
	}
	return true;
}

bool TCPSocket::doWrite()
{
	if (m_scktClient == INVALID_SOCKET) return false;
	if (m_scktState != Socket_Connected) return false;

	if (m_nOutbufLen == 0) return true;
	int send_res = send(m_scktClient, m_bufOutput, m_nOutbufLen, 0);
	if (send_res == SOCKET_ERROR) return false;
	if (send_res > m_nOutbufLen) return false;
	if (send_res == 0) return true;
	memmove(m_bufOutput, &m_bufOutput[send_res], m_nOutbufLen-send_res);  // move to top
	m_nOutbufLen -= send_res;
	return true;
}

void TCPSocket::callRemoteCommand(unsigned int cmd_id, va_list* pargs)
{
	if (!m_pProtocol) return;
	int pack_len = m_pProtocol->makePacket(m_bufMakePacket, OUTBUFSIZE,	cmd_id, pargs);
	if (pack_len <= 0) return;
	m_nMakePacketLen = pack_len;
	if (m_bFirstEncode)
	{
		memmove(&m_bufMakePacket[2], &m_bufMakePacket[0], m_nMakePacketLen);
		m_nMakePacketLen += 2;
		m_keyEncrypt = (unsigned char)((rand() % 255) + 1);
		m_idxEncryptTable = m_keyEncrypt % 10;
		m_bufMakePacket[0] = m_keyEncrypt;
		m_keyEncrypt = Encrypt(m_idxEncryptTable, m_keyEncrypt, (unsigned char*)&m_bufMakePacket[1],
			(unsigned char*)&m_bufMakePacket[1], m_nMakePacketLen - 1);
		m_bFirstEncode = false;
	}
	else
	{
		m_keyEncrypt = Encrypt(m_idxEncryptTable, m_keyEncrypt, (unsigned char*)m_bufMakePacket, 
			(unsigned char*)m_bufMakePacket, m_nMakePacketLen);
	}
	memcpy(&m_bufOutput[m_nOutbufLen], m_bufMakePacket, m_nMakePacketLen);
	m_nOutbufLen += m_nMakePacketLen;
}

void TCPSocket::destroy(void)
{
	// ??
	struct linger so_linger;
	so_linger.l_onoff = 1;
	so_linger.l_linger = 500;
	int ret = setsockopt(m_scktClient, SOL_SOCKET, SO_LINGER, (const char*)&so_linger, sizeof(so_linger));

	closeSocket();

	m_scktClient = INVALID_SOCKET;
	m_nInbufLen = 0;
	m_nOutbufLen = 0;

	memset(m_bufOutput, 0, sizeof(m_bufOutput));
	memset(m_bufInput, 0, sizeof(m_bufInput));
}

